#include <iostream>
#include "arm_neon.h"
#include <opencv2/opencv.hpp>

void matrix_multiply_c(float32_t *A, float32_t *B, float32_t *C, uint32_t n, uint32_t m, uint32_t k) {
    for (int i_idx=0; i_idx < n; i_idx++) {
        for (int j_idx=0; j_idx < m; j_idx++) {
            C[m*i_idx + j_idx] = 0;
            for (int k_idx=0; k_idx < k; k_idx++) {
                C[m*i_idx + j_idx] += A[k*i_idx + k_idx] * B[m * k_idx + j_idx];
            }
//             std::cout << C[m*i_idx + j_idx] << " ";
        }
    }
}


void matrix_multiply_neon(float32_t  *A, float32_t  *B, float32_t *C, uint32_t n, uint32_t m, uint32_t k) {
	/* 
	 * Multiply matrices A and B, store the result in C. 
	 * It is the user's responsibility to make sure the matrices are compatible.
	 */	

	int A_idx;
	int B_idx;
	int C_idx;
	
	// these are the columns of a 4x4 sub matrix of A
	float32x4_t A0;
	float32x4_t A1;
	float32x4_t A2;
	float32x4_t A3;
	
	// these are the columns of a 4x4 sub matrix of B
	float32x4_t B0;
	float32x4_t B1;
	float32x4_t B2;
	float32x4_t B3;
	
	// these are the columns of a 4x4 sub matrix of C
	float32x4_t C0;
	float32x4_t C1;
	float32x4_t C2;
	float32x4_t C3;
	
	for (int i_idx=0; i_idx<n; i_idx+=4) {
		for (int j_idx=0; j_idx<m; j_idx+=4) {
			// Zero accumulators before matrix op
			C0 = vmovq_n_f32(0);
			C1 = vmovq_n_f32(0);
			C2 = vmovq_n_f32(0);
			C3 = vmovq_n_f32(0);
			for (int k_idx=0; k_idx<k; k_idx+=4) {
				// Compute base index to 4x4 block
				A_idx = i_idx + n*k_idx;
				B_idx = k*j_idx + k_idx;
				
				// Load most current A values in row 
				A0 = vld1q_f32(A+A_idx);
				A1 = vld1q_f32(A+A_idx+n);
				A2 = vld1q_f32(A+A_idx+2*n);
				A3 = vld1q_f32(A+A_idx+3*n);
				
				// Multiply accumulate in 4x1 blocks, i.e. each column in C
				B0 = vld1q_f32(B+B_idx);
				C0 = vfmaq_laneq_f32(C0, A0, B0, 0);
				C0 = vfmaq_laneq_f32(C0, A1, B0, 1);
				C0 = vfmaq_laneq_f32(C0, A2, B0, 2);
				C0 = vfmaq_laneq_f32(C0, A3, B0, 3);
				
				B1 = vld1q_f32(B+B_idx+k);
				C1 = vfmaq_laneq_f32(C1, A0, B1, 0);
				C1 = vfmaq_laneq_f32(C1, A1, B1, 1);
				C1 = vfmaq_laneq_f32(C1, A2, B1, 2);
				C1 = vfmaq_laneq_f32(C1, A3, B1, 3);
				
				B2 = vld1q_f32(B+B_idx+2*k);
				C2 = vfmaq_laneq_f32(C2, A0, B2, 0);
				C2 = vfmaq_laneq_f32(C2, A1, B2, 1);
				C2 = vfmaq_laneq_f32(C2, A2, B2, 2);
				C2 = vfmaq_laneq_f32(C2, A3, B2, 3);
				
				B3 = vld1q_f32(B+B_idx+3*k);
				C3 = vfmaq_laneq_f32(C3, A0, B3, 0);
				C3 = vfmaq_laneq_f32(C3, A1, B3, 1);
				C3 = vfmaq_laneq_f32(C3, A2, B3, 2);
				C3 = vfmaq_laneq_f32(C3, A3, B3, 3);
			}
			// Compute base index for stores
			C_idx = n*j_idx + i_idx;
			vst1q_f32(C+C_idx, C0);
			vst1q_f32(C+C_idx+n, C1);
			vst1q_f32(C+C_idx+2*n, C2);
			vst1q_f32(C+C_idx+3*n, C3);
		}
	}
}


void matrix_multiply_4x4_neon(float32_t *A, float32_t *B, float32_t *C) {
	// these are the columns A
	float32x4_t A0;
	float32x4_t A1;
	float32x4_t A2;
	float32x4_t A3;
	
	// these are the columns B
	float32x4_t B0;
	float32x4_t B1;
	float32x4_t B2;
	float32x4_t B3;
	
	// these are the columns C
	float32x4_t C0;
	float32x4_t C1;
	float32x4_t C2;
	float32x4_t C3;
	
	A0 = vld1q_f32(A);
	A1 = vld1q_f32(A+4);
	A2 = vld1q_f32(A+8);
	A3 = vld1q_f32(A+12);
	
	// Zero accumulators for C values
	C0 = vmovq_n_f32(0);
	C1 = vmovq_n_f32(0);
	C2 = vmovq_n_f32(0);
	C3 = vmovq_n_f32(0);
	
	// Multiply accumulate in 4x1 blocks, i.e. each column in C
	B0 = vld1q_f32(B);
	C0 = vfmaq_laneq_f32(C0, A0, B0, 0);
	C0 = vfmaq_laneq_f32(C0, A1, B0, 1);
	C0 = vfmaq_laneq_f32(C0, A2, B0, 2);
	C0 = vfmaq_laneq_f32(C0, A3, B0, 3);
	vst1q_f32(C, C0);
	
	B1 = vld1q_f32(B+4);
	C1 = vfmaq_laneq_f32(C1, A0, B1, 0);
	C1 = vfmaq_laneq_f32(C1, A1, B1, 1);
	C1 = vfmaq_laneq_f32(C1, A2, B1, 2);
	C1 = vfmaq_laneq_f32(C1, A3, B1, 3);
	vst1q_f32(C+4, C1);
	
	B2 = vld1q_f32(B+8);
	C2 = vfmaq_laneq_f32(C2, A0, B2, 0);
	C2 = vfmaq_laneq_f32(C2, A1, B2, 1);
	C2 = vfmaq_laneq_f32(C2, A2, B2, 2);
	C2 = vfmaq_laneq_f32(C2, A3, B2, 3);
	vst1q_f32(C+8, C2);
	
	B3 = vld1q_f32(B+12);
	C3 = vfmaq_laneq_f32(C3, A0, B3, 0);
	C3 = vfmaq_laneq_f32(C3, A1, B3, 1);
	C3 = vfmaq_laneq_f32(C3, A2, B3, 2);
	C3 = vfmaq_laneq_f32(C3, A3, B3, 3);
	vst1q_f32(C+12, C3);
}

/*
    float32x4_t sum = vdupq_n_f32(0); //sum={0,0,0,0}
    float _a[] = {1,2,3,4}, _b[] = {5,6,7,8} ;
    float32x4_t a = vld1q_f32(_a), b = vld1q_f32(_b);
    //sum1={5,10,15,20}
    float32x4_t sum1 = vfmaq_laneq_f32(sum, a, b, 0); 
    //sum2={5,10,15,20}+{6,12,18,24} = {11,22,33,44}
    float32x4_t sum2 = vfmaq_laneq_f32(sum1, a, b, 1); 
    //sum3={11,22,33,44}+{7,14,21,28} = {18,36,54,72}
    float32x4_t sum3 = vfmaq_laneq_f32(sum2, a, b, 2);

    float32x4_t r1 = vextq_f32(a,b,1); //r1={2,3,4,5}
    float32x4_t r2 = vextq_f32(a,b,2); //r2={3,4,5,6}
    float32x4_t r3 = vextq_f32(a,b,3); //r3={4,5,6,7}
*/

int main(void){
    uint32_t n = 4;
    uint32_t m = 4;
    uint32_t k = 4;
//     cv::Mat A = cv::Mat::ones(n, k, CV_32FC1);
    cv::Mat A = (cv::Mat_<float32_t>(n,k) << 1,2,3,4, 5,6,7,8, 2,3,4,5, 6,7,8,9);
    cv::Mat B = cv::Mat::ones(k, m, CV_32FC1);
    // 对于neon，矩阵需要按列存储
//     cv::Mat A1 = (cv::Mat_<float32_t>(n,k) << 1,5,2,6, 2,6,3,7, 3,7,4,8, 4,8,5,9);
    cv::Mat A1 = A.t();
//     cv::Mat B1 = cv::Mat::ones(m, k, CV_32FC1);
    cv::Mat B1 = B.t();
    
    cv::Mat C = cv::Mat::zeros(n, m, CV_32FC1);
    cv::Mat D = cv::Mat::zeros(n, m, CV_32FC1);
    cv::Mat E = cv::Mat::zeros(n, m, CV_32FC1);
    
    matrix_multiply_c((float32_t*)A.data, (float32_t*)B.data, (float32_t*)C.data, n, m, k);
    matrix_multiply_4x4_neon((float32_t*)A1.data, (float32_t*)B1.data, (float32_t*)D.data);
    matrix_multiply_neon((float32_t*)A1.data, (float32_t*)B1.data, (float32_t*)E.data, n, m, k);
    
    std::cout << A << std::endl;
    std::cout << B << std::endl;
    std::cout << C << std::endl;
    std::cout << D.t() << std::endl;
    std::cout << E.t() << std::endl;
    
    return 0;
}
